import datetime
import pytz
import os, sys, os.path as osp
import json

from huggingface_hub import snapshot_download, hf_hub_url, get_hf_file_metadata
from huggingface_hub.hf_api import HfApi
from hf_torrent.utils import run_command, FORMAT_NAME, enumerate_hf_repo


REPO_BASE_DIR = "hf-repository"
TORRENT_BASE_DIR = "hf-torrent-store"
FORMAT_NAME = lambda s: s.replace("-", "_").replace("/", "-")


def main(repo="bert-base-uncased", delete_existing=False, overwrite=False):
    # ==================== Check whether generated ====================
    meta_info_fpath = osp.join(TORRENT_BASE_DIR, repo, "_hf_mirror_torrent.json")
    print(meta_info_fpath)
    if osp.exists(meta_info_fpath):
        print(
            "Metainfo file already exists. Checking whether it is generated by the latest commit."
        )
        with open(meta_info_fpath, "r") as fp:
            fpath_mapping = json.load(fp)

        api = HfApi()
        git_hash = api.repo_info(repo_id=repo).sha
        if (
            "lastest-commit" in fpath_mapping
            and fpath_mapping["lastest-commit"] == git_hash
            and not overwrite
        ):
            print(f"{repo} => {meta_info_fpath} already generated.")
            return

    # ==================== Download model ====================
    try:
        model_fpath = snapshot_download(repo)
    except Exception as e:
        print(e)
        print("Failed to download model. Skipping.")
        return

    hf_cache_base = osp.dirname(osp.dirname(model_fpath))

    fpath_mapping = {}
    fpath_mapping["fpath2uuid"] = {}
    fpath_mapping["uuid2fpath"] = {}

    # Create torrent for folder
    print("--" * 50)
    git_hash = osp.basename(model_fpath)
    torrent_path = osp.join(TORRENT_BASE_DIR, repo, f"_all.torrent")
    os.makedirs(osp.dirname(torrent_path), exist_ok=True)
    repo_name = FORMAT_NAME(repo)
    cmd = f"python py3createtorrent.py -t best5 {model_fpath} \
        --name '{git_hash}' \
        --webseed https://huggingface.co/{repo}/resolve/ \
        --webseed https://hf-mirror.com/{repo}/resolve/ \
        --output {torrent_path} --force"
    stdout, stderr = run_command(cmd)
    print(stdout, stderr)
    print(cmd)
    print("--" * 50)

    # Create torrent for folder (hf-mirror only, but the naming is more reable.)
    print("--" * 50)
    git_hash = osp.basename(model_fpath)
    torrent_path = osp.join(TORRENT_BASE_DIR, repo, f"_all_hf-mirror.torrent")
    os.makedirs(osp.dirname(torrent_path), exist_ok=True)

    def convert_repo_name(repo):
        if "/" in repo:
            org, name = repo.split("/")
            return f"{org}--{name}"
        else:
            return repo

    repo_folder = f"models--{convert_repo_name(repo)}--{git_hash[:7]}"  # 7 digits to follow HF conventional length.
    cmd = f"python py3createtorrent.py -t best5 {model_fpath} \
        --name '{repo_folder}' \
        --webseed https://hf-mirror.com/ws \
        --output {torrent_path} --force"
    stdout, stderr = run_command(cmd)
    print(stdout, stderr)
    print(cmd)
    print("--" * 50)

    # Create torrents for single-file
    for fpath in enumerate_hf_repo(model_fpath):
        print("--" * 50)
        file_name = osp.relpath(fpath, model_fpath)
        etag_hash = osp.basename(osp.realpath(fpath))
        repo_name = FORMAT_NAME(repo)
        torrent_name = FORMAT_NAME(file_name)

        print(repo_name, "\t", torrent_name, "\t", etag_hash)

        uuid = f"{repo_name}-{torrent_name}-{etag_hash}"

        torrent_path = osp.join(TORRENT_BASE_DIR, repo, f"{uuid}.torrent")
        os.makedirs(osp.dirname(torrent_path), exist_ok=True)

        hf_meta = get_hf_file_metadata(hf_hub_url(repo_id=repo, filename=file_name))
        commit_hash = hf_meta.commit_hash

        rel_fpath = osp.relpath(fpath, model_fpath)
        fpath_mapping["fpath2uuid"][rel_fpath] = f"{uuid}"
        fpath_mapping["uuid2fpath"][f"{uuid}"] = rel_fpath

        # if osp.exists(torrent_path):
        #     print(f"Skipping {torrent_path} as it already exists.")
        #     print("--" * 50)
        #     continue

        cmd = f"python py3createtorrent.py -t best5 {fpath} \
                --name '{uuid}' \
                --webseed https://huggingface.co/{repo}/resolve/{commit_hash}/{file_name} \
                --webseed https://hf-mirror.com/{repo}/resolve/{commit_hash}/{file_name} \
                --output {torrent_path} --force"
        stdout, stderr = run_command(cmd)
        print(stdout, stderr)
        print(cmd)
        print("--" * 50)

    desired_timezone = pytz.timezone("Asia/Shanghai")
    fpath_mapping["lastest-generated"] = datetime.datetime.now(
        desired_timezone
    ).strftime("%Y-%m-%d %H:%M:%S")
    fpath_mapping["lastest-commit"] = git_hash

    with open(osp.join(TORRENT_BASE_DIR, repo, "_hf_mirror_torrent.json"), "w") as fp:
        json.dump(fpath_mapping, fp, indent=2)

    if delete_existing:
        import shutil

        print("Removing cache: ", hf_cache_base)
        shutil.rmtree(hf_cache_base)


if __name__ == "__main__":
    import argparse

    # parser = argparse.ArgumentParser(prog='HF Torrent Creator')
    # parser.add_argument('repo')       # positional argument
    # args = parser.parse_args()
    main(overwrite=False)
